#include <stdio.h>
#include <stdlib.h>

#include <libavcodec/avcodec.h>
#include <libavformat/avformat.h>
#include <libavutil/imgutils.h>
#include <libswscale/swscale.h>
#include <libavdevice/avdevice.h>

int AVFrame2JPG(AVFrame *pfrm, enum AVPixelFormat pix_fmt, int des_width, int des_height, char *jpgname) {

    AVFormatContext *pfmt_ctx = NULL;
    AVStream *pstm = NULL;
    AVCodecParameters *pcodecpar = NULL;
    AVPacket packet;
    AVCodecContext *pcodec_ctx = NULL;
    AVCodec *pcodec = NULL;
    struct SwsContext *psws_ctx = NULL;
    AVFrame *pyuv_frm = NULL;

    int ret = 0;
    int width = pfrm->width;
    int height = pfrm->height;

    if ((pyuv_frm = av_frame_alloc()) == NULL) {
        av_log(NULL, AV_LOG_ERROR, "av_frame_alloc ERROR \n");
        return AVERROR(ENOMEM);
    }

    if ((psws_ctx = sws_getContext(width, height, pix_fmt, des_width, des_height, AV_PIX_FMT_YUVJ422P, SWS_FAST_BILINEAR, NULL, NULL, NULL)) == NULL) {
        av_log(NULL, AV_LOG_ERROR, "sws_getContext ERROR \n");
        return AVERROR(EINVAL);
    }

    if ((ret = av_image_alloc(pyuv_frm->data, pyuv_frm->linesize, des_width, des_height, AV_PIX_FMT_YUVJ422P, 1)) < 0) {
        av_log(NULL, AV_LOG_ERROR, "av_image_alloc ERROR \n");
        return AVERROR(EINVAL);
    }
    
    pyuv_frm->width = des_width;
    pyuv_frm->height = des_height;
    pyuv_frm->format = AV_PIX_FMT_YUVJ422P;

    sws_scale(psws_ctx, (const uint8_t * const *) pfrm->data, pfrm->linesize, 0, height, pyuv_frm->data, pyuv_frm->linesize);

    ret = avformat_alloc_output_context2(&pfmt_ctx, NULL, "mjpeg", jpgname);
    if (ret < 0) {
        av_log(NULL, AV_LOG_ERROR, "avformat_alloc_context ERROR %s in AVFrame2JPG \n", av_err2str(ret));
        return AVERROR(ret);
    }   

    if ((ret = avio_open2(&(pfmt_ctx->pb), jpgname, AVIO_FLAG_WRITE, NULL, NULL)) < 0) {
        av_log(NULL, AV_LOG_ERROR, "avio_open2 ERROR %s in AVFrame2JPG \n", av_err2str(ret));
        return AVERROR(ret);
    }

    pstm = avformat_new_stream(pfmt_ctx, NULL);
    if (pstm == NULL) {
        av_log(NULL, AV_LOG_ERROR, "avformat_new_stream ERROR in AVFrame2JPG \n");
        return AVERROR(ENOMEM);
    }

    pcodecpar = pstm->codecpar;
    pcodecpar->height = des_height;
    pcodecpar->width = des_width;
    pcodecpar->codec_id = pfmt_ctx->oformat->video_codec;
    pcodecpar->codec_type = AVMEDIA_TYPE_VIDEO;
    pcodecpar->format = AV_PIX_FMT_YUVJ422P;

    pcodec = avcodec_find_encoder(pcodecpar->codec_id);
    if (pcodec == NULL) {
        av_log(NULL, AV_LOG_ERROR, "avcodec_find_encoder ERROR in AVFrame2JPG \n");
        return AVERROR(EINVAL);
    }
    //printf("codec name: %s, long name: %s \n", pcodec->name, pcodec->long_name);

    pcodec_ctx = avcodec_alloc_context3(pcodec);
    if (pcodec_ctx == NULL) {
        av_log(NULL, AV_LOG_ERROR, "avcodec_alloc_context3 ERROR in main \n");
        return AVERROR(ENOMEM);
    }

    ret = avcodec_parameters_to_context(pcodec_ctx, pstm->codecpar);
    if (ret < 0) {
        av_log(NULL, AV_LOG_ERROR, "avcodec_parameters_to_context ERROR %s in AVFrame2JPG \n", av_err2str(ret));
        return AVERROR(ret);
    }

    pcodec_ctx->time_base.den = 25;
    pcodec_ctx->time_base.num = 1;

    //** change the quality of jpg
    pcodec_ctx->qcompress = 80.0/100.0;
    pcodec_ctx->qmin = 2;
    pcodec_ctx->qmax = 52;
    pcodec_ctx->max_qdiff = 5;

    if ((ret = avcodec_open2(pcodec_ctx, pcodec, NULL)) < 0){
        av_log(NULL, AV_LOG_ERROR, "avcodec_open2 ERROR %s in AVFrame2JPG \n", av_err2str(ret));
        return AVERROR(ret);
    }

    ret = avformat_write_header(pfmt_ctx, NULL);
    if (ret < 0) {
        av_log(NULL, AV_LOG_ERROR, "avformat_write_header ERROR %s in AVFrame2JPG \n", av_err2str(ret));
        return AVERROR(ret);
    }
    
    ret = av_new_packet(&packet, 3 * des_width * des_height);
    if (ret < 0) {
        av_log(NULL, AV_LOG_ERROR, "av_new_packet ERROR in AVFrame2JPG \n");
        return -1;
    }

    ret = avcodec_send_frame(pcodec_ctx, pyuv_frm);
    if (ret < 0) {
        av_log(NULL, AV_LOG_ERROR, "avcodec_send_frame ERROR %s in AVFrame2JPG \n", av_err2str(ret));
        return AVERROR(ret);
    }

    while (ret >= 0) {
        ret = avcodec_receive_packet(pcodec_ctx, &packet);
        if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF ) {
            break;
        } else if (ret < 0) {
            av_log(NULL, AV_LOG_ERROR, "avcodec_receive_packet ERROR %s in AVFrame2JPG \n", av_err2str(ret));
            return AVERROR(ret);
        }
        
        if ((ret = av_write_frame(pfmt_ctx, &packet)) < 0) {
            av_log(NULL, AV_LOG_ERROR, "av_write_frame ERROR %s in AVFrame2JPG \n", av_err2str(ret));
            return AVERROR(ret);
        }
        av_packet_unref(&packet);
        av_write_trailer(pfmt_ctx);
    }

    av_frame_free(&pyuv_frm);
    avcodec_close(pcodec_ctx);
    avformat_free_context(pfmt_ctx);
    avio_close(pfmt_ctx->pb);
    sws_freeContext(psws_ctx);

    return 0;
}


int main(int argc, char* argv[]){

    int ret = 0;

    AVFormatContext *pfmt_ctx = NULL;
    AVInputFormat *infmt = NULL;
    AVCodecParameters *pcodec_par = NULL;
    AVCodecContext *pcodec_ctx = NULL;
    AVCodec *pcodec = NULL;
    AVPacket *ppkt = NULL;
    AVFrame *pfrm = NULL;
    char outname[100];
    enum AVPixelFormat pix_fmt;

    int video_idx = -1;
    int frame_idx = 0;
    int des_width = 0;
    int des_height = 0;

    avdevice_register_all();
    
    if ((pfmt_ctx = avformat_alloc_context()) == NULL) {
        av_log(NULL, AV_LOG_ERROR, "avformat_alloc_context ERROR \n");
        return AVERROR(ENOMEM);
    }

    if ((infmt = av_find_input_format("video4linux2")) == NULL) {
        av_log(NULL, AV_LOG_ERROR, "av_find_input_format ERROR \n");
        return AVERROR(EINVAL);
    }

    if((ret = avformat_open_input(&pfmt_ctx, "/dev/video0", infmt, NULL)) < 0) {
        av_log(NULL, AV_LOG_ERROR, "avformat_open_input ERROR \n");
        return AVERROR(EINVAL);
    }

    if ((ret = avformat_find_stream_info(pfmt_ctx, NULL)) < 0) {
        av_log(NULL, AV_LOG_ERROR, "avformat_find_stream_info ERROR \n");
        return AVERROR(EINVAL);
    }

    for (int i=0; i<pfmt_ctx->nb_streams; i++) {
        if (pfmt_ctx->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO) {
            video_idx = i;
            break;
        }
    }

    if (video_idx == -1) {
        fprintf(stderr, "Can not find video index \n");
        return -1;
    }

    pcodec_par = pfmt_ctx->streams[video_idx]->codecpar;

    if ((pcodec = avcodec_find_decoder(pcodec_par->codec_id)) == NULL) {
        av_log(NULL, AV_LOG_ERROR, "avcodec_find_encoder ERROR \n");
        return AVERROR(ret);
    }
    
    if ((pcodec_ctx = avcodec_alloc_context3(pcodec)) == NULL) {
        av_log(NULL, AV_LOG_ERROR, "avcodec_alloc_context3 ERROR \n");
        return AVERROR(ENOMEM);
    }

    if ((ret = avcodec_parameters_to_context(pcodec_ctx, pcodec_par)) < 0) {
        av_log(NULL, AV_LOG_ERROR, "avcodec_parameters_to_context ERROR \n");
        return AVERROR(ret);
    }
    
    pcodec_ctx->time_base.den = 25;
    pcodec_ctx->time_base.num = 1;
    pix_fmt = pcodec_ctx->pix_fmt;
    
    //des_width = pcodec_ctx->width;
    //des_height = pcodec_ctx->height;
    
    des_width = 1280;
    des_height = 720;

    if ((ret = avcodec_open2(pcodec_ctx, pcodec, NULL)) < 0) {
        av_log(NULL, AV_LOG_ERROR, "avcodec_open2 ERROR \n");
        return AVERROR(ret);
    }

    if ((ppkt = av_packet_alloc()) == NULL) {
        av_log(NULL, AV_LOG_ERROR, "av_packet_alloc ERROR \n");
        return AVERROR(ENOMEM);
    }

    if ((pfrm = av_frame_alloc()) == NULL) {
        av_log(NULL, AV_LOG_ERROR, "av_frame_alloc ERROR \n");
        return AVERROR(ENOMEM);
    }

    while ((ret = av_read_frame(pfmt_ctx, ppkt)) >= 0) {
        if (ppkt->stream_index == video_idx) {
            if ((ret = avcodec_send_packet(pcodec_ctx, ppkt)) < 0) {
                av_log(NULL, AV_LOG_ERROR, "avcodec_send_packet ERROR %s \n", av_err2str(ret));
                return AVERROR(ret);
            }

            while (ret >= 0) {
                ret = avcodec_receive_frame(pcodec_ctx, pfrm);
                if (ret == AVERROR_EOF || ret == AVERROR(EAGAIN)) {
                    break;
                } else if (ret < 0) {
                    av_log(NULL, AV_LOG_ERROR, "avcodec_receive_frame ERROR \n");
                    return AVERROR(ret);
                }

                frame_idx++;
                sprintf(outname, "%d.jpg", frame_idx);
                fprintf(stdout, "Output file name : %s \n", outname);

                if ((ret = AVFrame2JPG(pfrm, pix_fmt, des_width, des_height, outname)) < 0) {
                    av_log(NULL, AV_LOG_ERROR, "AVFrame2JPG ERROR %s \n", av_err2str(ret));
                    return AVERROR(ret);
                }

                av_frame_unref(pfrm);
            }
        }
        av_packet_unref(ppkt);
        if (frame_idx == 100) {
            break;
        }
    }

    avformat_close_input(&pfmt_ctx);
    avformat_free_context(pfmt_ctx);
    av_frame_free(&pfrm);
    av_packet_free(&ppkt);

    return 0;
}